{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# GAMENet Model Training on MIMIC-III Dataset\n",
        "\n",
        "Train the GAMENet (Graph Augmented MEmory Networks) model for medication recommendation using the MIMIC-III dataset.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/Caskroom/miniforge/base/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "No config path provided, using default config\n",
            "Initializing mimic3 dataset from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III (dev mode: True)\n",
            "Scanning table: patients from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv\n",
            "Some column names were converted to lowercase\n",
            "Scanning table: admissions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Some column names were converted to lowercase\n",
            "Scanning table: icustays from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv\n",
            "Some column names were converted to lowercase\n",
            "Scanning table: diagnoses_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv\n",
            "Some column names were converted to lowercase\n",
            "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Scanning table: procedures_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv\n",
            "Some column names were converted to lowercase\n",
            "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Scanning table: prescriptions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv\n",
            "Some column names were converted to lowercase\n",
            "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Collecting global event dataframe...\n",
            "Dev mode enabled: limiting to 1000 patients\n",
            "Collected dataframe with shape: (46835, 49)\n",
            "Dataset: mimic3\n",
            "Dev mode: True\n",
            "Number of patients: 1000\n",
            "Number of events: 46835\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.datasets import MIMIC3Dataset\n",
        "\n",
        "dataset = MIMIC3Dataset(\n",
        "    root=\"https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III\",\n",
        "    tables=[\"DIAGNOSES_ICD\", \"PROCEDURES_ICD\", \"PRESCRIPTIONS\"],\n",
        "    dev=True,\n",
        ")\n",
        "dataset.stats()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Set Drug Recommendation Task\n",
        "\n",
        "We use the drug recommendation task which predicts medications based on patient conditions and procedures.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Setting task DrugRecommendationMIMIC3 for mimic3 base dataset...\n",
            "Generating samples with 1 worker(s)...\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Generating samples for DrugRecommendationMIMIC3 with 1 worker: 100%|██████████| 1000/1000 [00:01<00:00, 546.61it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Label drugs vocab: {'*NF*': 0, '0.9%': 1, '1/2 ': 2, '5% D': 3, 'ACD-': 4, 'Acet': 5, 'Acyc': 6, 'Aden': 7, 'Albu': 8, 'Alen': 9, 'Allo': 10, 'Alpr': 11, 'Alte': 12, 'Alum': 13, 'Ambi': 14, 'Amik': 15, 'Amin': 16, 'Amio': 17, 'Amlo': 18, 'Amph': 19, 'Ampi': 20, 'Apro': 21, 'Aqua': 22, 'Arti': 23, 'Asco': 24, 'Aspi': 25, 'Aten': 26, 'Ator': 27, 'Atro': 28, 'Azit': 29, 'Becl': 30, 'Bell': 31, 'Bici': 32, 'Bisa': 33, 'Bism': 34, 'Brom': 35, 'BuPR': 36, 'Bume': 37, 'BusP': 38, 'Calc': 39, 'Capt': 40, 'Carb': 41, 'Carv': 42, 'Casp': 43, 'Cefa': 44, 'Cefe': 45, 'Ceft': 46, 'Cele': 47, 'Ceph': 48, 'Cety': 49, 'Chlo': 50, 'Cipr': 51, 'Cisa': 52, 'Cita': 53, 'Citr': 54, 'Clin': 55, 'Clob': 56, 'Clon': 57, 'Clop': 58, 'Clot': 59, 'Colc': 60, 'Coll': 61, 'Cosy': 62, 'Cyan': 63, 'D5 1': 64, 'D5NS': 65, 'D5W': 66, 'D5W ': 67, 'DOBU': 68, 'Daki': 69, 'Dapt': 70, 'Desm': 71, 'Dexa': 72, 'Dexm': 73, 'Dext': 74, 'Diaz': 75, 'Dida': 76, 'Digo': 77, 'Dilt': 78, 'Diov': 79, 'Diph': 80, 'Diva': 81, 'Dobu': 82, 'Docu': 83, 'Dola': 84, 'DopA': 85, 'Dorz': 86, 'Doxa': 87, 'Dron': 88, 'Efav': 89, 'Emtr': 90, 'Enal': 91, 'Enox': 92, 'Epin': 93, 'Eple': 94, 'Epoe': 95, 'Epti': 96, 'Esmo': 97, 'Etom': 98, 'Euce': 99, 'Ezet': 100, 'Famo': 101, 'Fent': 102, 'Ferr': 103, 'Fexo': 104, 'Fina': 105, 'Flee': 106, 'Fluc': 107, 'Flud': 108, 'Flut': 109, 'Foli': 110, 'Fond': 111, 'Fosa': 112, 'Furo': 113, 'Gaba': 114, 'Gast': 115, 'Gent': 116, 'Glip': 117, 'Gluc': 118, 'GlyB': 119, 'Glyb': 120, 'Glyc': 121, 'Goly': 122, 'Guai': 123, 'HYDR': 124, 'Halo': 125, 'Hepa': 126, 'Humu': 127, 'Hydr': 128, 'Ibup': 129, 'Idar': 130, 'Imip': 131, 'Infl': 132, 'Insu': 133, 'Ipra': 134, 'Isos': 135, 'Keto': 136, 'LR': 137, 'Labe': 138, 'Lact': 139, 'Lans': 140, 'Leva': 141, 'Leve': 142, 'Levo': 143, 'Lido': 144, 'Line': 145, 'Lisi': 146, 'Lope': 147, 'Lora': 148, 'Lyri': 149, 'Magn': 150, 'Mecl': 151, 'Mege': 152, 'Mepe': 153, 'Mero': 154, 'MetF': 155, 'MetR': 156, 'Metf': 157, 'Meth': 158, 'Meto': 159, 'Metr': 160, 'Mico': 161, 'Mida': 162, 'Mido': 163, 'Milk': 164, 'Milr': 165, 'Mine': 166, 'Mirt': 167, 'Mont': 168, 'Morp': 169, 'Movi': 170, 'Mult': 171, 'Mupi': 172, 'Myco': 173, 'NIFE': 174, 'NORe': 175, 'NS': 176, 'NS (': 177, 'Nado': 178, 'Nafc': 179, 'Neos': 180, 'Neph': 181, 'Nesi': 182, 'Neut': 183, 'Nico': 184, 'Nitr': 185, 'Nore': 186, 'Nyst': 187, 'Olan': 188, 'Omep': 189, 'Onda': 190, 'Opiu': 191, 'Oxca': 192, 'Oxyc': 193, 'Oxym': 194, 'Pant': 195, 'Papa': 196, 'Phen': 197, 'Phyt': 198, 'Piog': 199, 'Pipe': 200, 'Pneu': 201, 'Poly': 202, 'Pota': 203, 'Prav': 204, 'Pred': 205, 'Proc': 206, 'Prom': 207, 'Prop': 208, 'Prot': 209, 'Quet': 210, 'Qvar': 211, 'Ralt': 212, 'Rani': 213, 'Read': 214, 'Ropi': 215, 'Rosi': 216, 'Sarn': 217, 'Scop': 218, 'Senn': 219, 'Sert': 220, 'Seve': 221, 'Sime': 222, 'Simv': 223, 'Sodi': 224, 'Spir': 225, 'Succ': 226, 'Sucr': 227, 'Sulf': 228, 'Tacr': 229, 'Tams': 230, 'Tema': 231, 'Thia': 232, 'Timo': 233, 'Tiot': 234, 'Tiza': 235, 'TraM': 236, 'Trav': 237, 'Tric': 238, 'Tube': 239, 'Vals': 240, 'Vanc': 241, 'Vaso': 242, 'Vecu': 243, 'Venl': 244, 'Vita': 245, 'Vori': 246, 'Warf': 247, 'Zinc': 248, 'Zolp': 249, 'cilo': 250, 'traM': 251, 'traZ': 252}\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Processing samples: 100%|██████████| 51/51 [00:00<00:00, 2501.75it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generated 51 samples for task DrugRecommendationMIMIC3\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.tasks import DrugRecommendationMIMIC3\n",
        "\n",
        "task = DrugRecommendationMIMIC3()\n",
        "samples = dataset.set_task(task)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Split Dataset\n",
        "\n",
        "Split the dataset into train, validation, and test sets using patient-level splitting.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "from pyhealth.datasets import split_by_patient, get_dataloader\n",
        "\n",
        "train_dataset, val_dataset, test_dataset = split_by_patient(\n",
        "    samples, ratios=[0.7, 0.15, 0.15]\n",
        ")\n",
        "\n",
        "train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True)\n",
        "val_loader = get_dataloader(val_dataset, batch_size=64, shuffle=False)\n",
        "test_loader = get_dataloader(test_dataset, batch_size=64, shuffle=False)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialize GAMENet Model\n",
        "\n",
        "Create the GAMENet model with specified hyperparameters.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/arjunchatterjee/PyHealth/pyhealth/sampler/sage_sampler.py:3: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
            "  import pkg_resources\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.models import GAMENet\n",
        "\n",
        "model = GAMENet(\n",
        "    dataset=samples,\n",
        "    embedding_dim=128,\n",
        "    hidden_dim=128,\n",
        "    num_layers=1,\n",
        "    dropout=0.5,\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Train Model\n",
        "\n",
        "Train the model using the PyHealth Trainer with relevant metrics for drug recommendation.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GAMENet(\n",
            "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
            "    (conditions): Embedding(303, 128)\n",
            "    (procedures): Embedding(115, 128)\n",
            "    (drugs_hist): Embedding(175, 128)\n",
            "  ))\n",
            "  (cond_rnn): GRU(128, 128, batch_first=True)\n",
            "  (proc_rnn): GRU(128, 128, batch_first=True)\n",
            "  (query): Sequential(\n",
            "    (0): ReLU()\n",
            "    (1): Linear(in_features=256, out_features=128, bias=True)\n",
            "  )\n",
            "  (gamenet): GAMENetLayer(\n",
            "    (ehr_gcn): GCN(\n",
            "      (gcn1): GCNLayer()\n",
            "      (dropout_layer): Dropout(p=0.5, inplace=False)\n",
            "      (gcn2): GCNLayer()\n",
            "    )\n",
            "    (ddi_gcn): GCN(\n",
            "      (gcn1): GCNLayer()\n",
            "      (dropout_layer): Dropout(p=0.5, inplace=False)\n",
            "      (gcn2): GCNLayer()\n",
            "    )\n",
            "    (fc): Linear(in_features=384, out_features=253, bias=True)\n",
            "    (bce_loss_fn): BCEWithLogitsLoss()\n",
            "  )\n",
            ")\n",
            "Metrics: ['jaccard_samples', 'f1_samples', 'pr_auc_samples', 'ddi']\n",
            "Device: cpu\n",
            "\n",
            "Training:\n",
            "Batch size: 64\n",
            "Optimizer: <class 'torch.optim.adam.Adam'>\n",
            "Optimizer params: {'lr': 0.001}\n",
            "Weight decay: 0.0\n",
            "Max grad norm: None\n",
            "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x13cabf980>\n",
            "Monitor: jaccard_samples\n",
            "Monitor criterion: max\n",
            "Epochs: 5\n",
            "Patience: None\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 0 / 5: 100%|██████████| 1/1 [00:00<00:00, 13.48it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-0, step-1 ---\n",
            "loss: 0.6989\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 137.55it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-0, step-1 ---\n",
            "jaccard_samples: 0.0450\n",
            "f1_samples: 0.0816\n",
            "pr_auc_samples: 0.0664\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6768\n",
            "New best jaccard_samples score (0.0450) at epoch-0, step-1\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 1 / 5: 100%|██████████| 1/1 [00:00<00:00, 27.67it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-1, step-2 ---\n",
            "loss: 0.6772\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 141.89it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-1, step-2 ---\n",
            "jaccard_samples: 0.0450\n",
            "f1_samples: 0.0816\n",
            "pr_auc_samples: 0.1666\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6528\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 2 / 5: 100%|██████████| 1/1 [00:00<00:00, 27.42it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-2, step-3 ---\n",
            "loss: 0.6552\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 93.62it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-2, step-3 ---\n",
            "jaccard_samples: 0.0450\n",
            "f1_samples: 0.0816\n",
            "pr_auc_samples: 0.2112\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6264\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 3 / 5: 100%|██████████| 1/1 [00:00<00:00, 24.14it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-3, step-4 ---\n",
            "loss: 0.6312\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 100.81it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-3, step-4 ---\n",
            "jaccard_samples: 0.0450\n",
            "f1_samples: 0.0816\n",
            "pr_auc_samples: 0.2289\n",
            "ddi_score: 0.0000\n",
            "loss: 0.5945\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 4 / 5: 100%|██████████| 1/1 [00:00<00:00, 25.92it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-4, step-5 ---\n",
            "loss: 0.6022\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 132.22it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-4, step-5 ---\n",
            "jaccard_samples: 0.0450\n",
            "f1_samples: 0.0816\n",
            "pr_auc_samples: 0.2485\n",
            "ddi_score: 0.0000\n",
            "loss: 0.5580\n",
            "Loaded best model\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.trainer import Trainer\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    metrics=[\"jaccard_samples\", \"f1_samples\", \"pr_auc_samples\", \"ddi\"],\n",
        ")\n",
        "\n",
        "trainer.train(\n",
        "    train_dataloader=train_loader,\n",
        "    val_dataloader=val_loader,\n",
        "    epochs=5,\n",
        "    monitor=\"jaccard_samples\",\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluate on Test Set\n",
        "\n",
        "Evaluate the trained model on the test set and print the results.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 112.65it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Test Set Results:\n",
            "  Jaccard (samples): 0.1210\n",
            "  F1 (samples): 0.2091\n",
            "  PR-AUC (samples): 0.1661\n",
            "  DDI Rate: N/A (metric not available)\n"
          ]
        }
      ],
      "source": [
        "results = trainer.evaluate(test_loader)\n",
        "\n",
        "print(\"Test Set Results:\")\n",
        "print(f\"  Jaccard (samples): {results['jaccard_samples']:.4f}\")\n",
        "print(f\"  F1 (samples): {results['f1_samples']:.4f}\")\n",
        "print(f\"  PR-AUC (samples): {results['pr_auc_samples']:.4f}\")\n",
        "ddi_value = results.get(\"ddi\")\n",
        "if ddi_value is not None:\n",
        "    print(f\"  DDI Rate: {ddi_value:.4f}\")\n",
        "else:\n",
        "    print(\"  DDI Rate: N/A (metric not available)\")\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "base",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
